import asyncio
import base64
import json
import logging
import os

import cv2

import nanoid
from ultralytics import YOLO
from pathlib import Path
from board.TensorBoard import TBServer
from common import ResultUtils
from common.DataUtils import convert_mask_to_coco


class TaskManager:
    def __init__(self):

        self.training_tasks = {}
        self.tensorboard_tasks = {}
        self.data_converting_tasks = {}
        self.models = {}

    def create_training_task(self,params):
        task_id = params["pretrain"].split(".")[0]+"-"+nanoid.generate()
        params["name"] = task_id
        new_task = asyncio.create_task(self.training_model(params))
        self.training_tasks[task_id] = new_task
        return task_id


    async def training_model(self, params):
        pretrain = params["pretrain"]
        if pretrain.startswith("yolo"):
            params_for_train = {
                "data": params["coco_file"],
                "epochs": int(params["epochs"]),
                "batch": int(params["batch"]),
                "imgsz": int(params["imgsz"]),
                "device": params["device"],
                "project": params["train_log_path"],
                "name": params["name"],
                "weight_decay": float(params["weight_decay"]),
                "patience": int(params["patience"]),
                "dropout": float(params["dropout"])
            }
            pretrain_model_path = Path(params["pretrain_path"] + "/" + pretrain)
            if pretrain_model_path.exists():
                logging.info(f"starting training,using pretrain model:{pretrain_model_path}")
                model = YOLO(params["pretrain_path"] + "/" + pretrain)
                model.train(**params_for_train)


    def deploy_model(self, params):
        status = -1
        msg = ""
        model_path = params["model_path"]
        model_id = params["model_id"]
        if model_id.startswith("yolo"):
            model = YOLO(model_path)
            if model_id in self.models:
                status = -1
                msg = f"model {model_id} already existed."
            else:
                self.models[model_id] = model
                status = 0
                msg = f"model {model_id} has been deployed"
        return status, msg



    def cv_predict_image(self, params):
        model_id = params["model_id"]
        colors_path = params["colors"]
        task_type = params["task_type"]
        with open(colors_path, "r") as f:
            colors = json.load(f)
        img_data_64encode = params["image"]
        img_data_64decode = base64.b64decode(img_data_64encode)
        task_id = nanoid.generate()
        with open(f"{task_id}.png", "wb") as f:
            f.write(img_data_64decode)
        result = self.models[model_id].predict(f"{task_id}.png")[0]
        ori_img_data_cv = cv2.imread(f"{task_id}.png")
        ret = {}
        output_img_path = f"{task_id}_output.png"
        if task_type == "segmentation":
            ret = ResultUtils.seg_result_proc(ori_img_data_cv, output_img_path, result, colors)
        if task_type == "object_detection":
            ret = ResultUtils.obt_result_proc(ori_img_data_cv, output_img_path, result, colors)
        os.remove(f"{task_id}.png")
        return f"{task_id}_output.png", ret

    def create_convertMask2CoCo_task(self, classes, masks_dir, type="train", suffix=".png"):
        convert_mask_to_coco(classes, masks_dir, type, suffix)

    def convert_coco_to_yolo(self):
        print()

    def create_tensorboard_task(self, taskId, logDir, port):
        tbs = TBServer(logDir, port)
        process = tbs.start()
        self.tensorboard_tasks[taskId] = process

    def stop_tensorboard(self, taskId):
        process = self.tensorboard_tasks.get(taskId)
        if process:
            process.terminate()
            del self.tensorboard_tasks[taskId]
            return True
        else:
            return False
